import os
import numpy as np
import cv2
import random
from PIL import Image
import torchvision.transforms as transforms


def shuffle_split_data(data,testratio=0.2):
    # random.shuffle(real)
    # random.shuffle(fake)
    # len1 = len(real)
    # len2 = len(fake)
    # tlen1 = int(len1*testratio)
    # tlen2 = int(len2*testratio)
    
    # realtest,faketest = real[:tlen1],fake[:tlen2]
    # realtrain,faketrain = real[tlen1:],fake[tlen2:]
    
    random.shuffle(data)
    len1 = len(data)
    tlen1 = int(len1*testratio)
    
    datatest= data[:tlen1]
    datatrain= data[tlen1:]

    return datatrain,datatest

def searchalldata(root):
    domains = os.listdir(root)
    domains.remove('gendata.py')
    dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    for domain in domains:
        cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            s0paths = os.listdir(os.path.join(root,domain,cat,'S0'))
            dolppaths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
            s0paths.sort()
            dolppaths.sort()
            if 'geniune' in domain:
                for (s0path,dolppath) in zip(s0paths,dolppaths):
                    s0real.append((os.path.join(root,domain,cat,s0path),1))
                    dolpreal.append((os.path.join(root,domain,cat,dolppath),1))
            else:
                for (s0path,dolppath) in zip(s0paths,dolppaths):
                    s0fake.append((os.path.join(root,domain,cat,s0path),0))
                    dolpfake.append((os.path.join(root,domain,cat,dolppath),0))
    return dolpreal,dolpfake,s0real,s0fake

def split_real_fake(dolpreal,dolpfake,s0real,s0fake,testratio=0.2):
    if len(dolpreal) > 0 :
        indreal = np.arange(0,len(dolpreal))
        indrealtrain,indrealtest = shuffle_split_data(indreal,testratio)
        dolpreal,s0real = np.array(dolpreal),np.array(s0real)
        
    if len(dolpfake) > 0:
        indfake = np.arange(0,len(dolpfake))
        indfaketrain,indfaketest = shuffle_split_data(indfake,testratio)
        dolpfake,s0fake = np.array(dolpfake),np.array(s0fake)
    # indrealtrain,indfaketrain,indrealtest,indfaketest = shuffle_split_data(indreal,indfake,testratio)
    
    # dolpreal,dolpfake,s0real,s0fake = np.array(dolpreal),np.array(dolpfake),np.array(s0real),np.array(s0fake)
    dolptraindir,dolptestdir,s0traindir,s0testdir = [],[],[],[]
    if len(dolpreal) > 0 and len(dolpfake) > 0:
        # dolptraindir,dolptestdir = np.vstack((dolpreal[indrealtrain],dolpfake[indfaketrain])),np.vstack((dolpreal[indrealtest],dolpfake[indfaketest]))
        # s0traindir,s0testdir = np.vstack((s0real[indrealtrain],s0fake[indfaketrain])),np.vstack((s0real[indrealtest],s0fake[indfaketest]))
        dolptraindir,dolptestdir = list(dolpreal[indrealtrain])+list(dolpfake[indfaketrain]),list(dolpreal[indrealtest])+list(dolpfake[indfaketest])
        s0traindir,s0testdir = list(s0real[indrealtrain])+list(s0fake[indfaketrain]),list(s0real[indrealtest])+list(s0fake[indfaketest])
    elif len(dolpreal) > 0 and len(dolpfake) == 0:
        dolptraindir,dolptestdir = list(dolpreal[indrealtrain]),list(dolpreal[indrealtest])
        if len(s0real) == len(dolpreal):
            s0traindir,s0testdir = list(s0real[indrealtrain]),list(s0real[indrealtest])
    elif len(dolpreal) == 0 and len(dolpfake) > 0:
        dolptraindir,dolptestdir = list(dolpfake[indfaketrain]),list(dolpfake[indfaketest])
        if len(s0fake) == len(dolpfake):
            s0traindir,s0testdir = list(s0fake[indfaketrain]),list(s0fake[indfaketest])
    
    # dolptraindir+=_dolptraindir
    # dolptestdir+=dolptestdir
    # s0traindir+=s0traindir
    # s0testdir+=s0testdir
        
    return np.array(dolptraindir),np.array(dolptestdir),np.array(s0traindir),np.array(s0testdir)


def save_train_test_to_txt(traindir,testdir,txtsavepath,filenameprefix,version):
    # traindir,testdir=gendatalist(datadir)
    # with open(os.path.join(txtsavepath,filenameprefix+'all.txt'),'w') as f:
    #     all_ = traindir + testdir
    #     for line in all_:
    #         f.write(line[0]+','+str(line[1])+'\n')
            
    with open(os.path.join(txtsavepath,filenameprefix+'_'+version+'_train.txt'),'w') as f:
        for line in traindir:
            f.write(line.__str__()+'\n')
    
    with open(os.path.join(txtsavepath,filenameprefix+'_'+version+'_test.txt'),'w') as f:
        for line in testdir:
            f.write(line.__str__()+'\n')


def gendatalist(root,dis,testratio=0.2):
    # domains = ['fake']['S0']
    # domains.remove('gendata.py')
    domains = ['geniune_domain']
    cats = ['HUT']
    # dolpreal,dolpfake,s0real,s0fake = [],[],[],[]
    data_dict = {}
    data_dict['S0']={}
    data_dict['S1']={}
    data_dict['S2']={}
    data_dict['DOLP']={}
    
    
    for domain in domains:
        # cats = os.listdir(os.path.join(root,domain))
        for cat in cats:
            s0paths = os.listdir(os.path.join(root,domain,cat,'S0'))
            s1paths = os.listdir(os.path.join(root,domain,cat,'S1'))
            s2paths = os.listdir(os.path.join(root,domain,cat,'S2'))
            dolppaths = os.listdir(os.path.join(root,domain,cat,'DOLP'))
            
            s0paths.sort()
            s1paths.sort()
            s2paths.sort()
            dolppaths.sort()

            for (s0path,s1path,s2path,dolppath) in zip(s0paths,s1paths,s2paths,dolppaths):
                tmp = int(s0path.split('_')[0])
                dis_mode = s0path.split('_')[1]
                if dis_mode in dis:
                    if not tmp in data_dict['S0'].keys():
                        data_dict['S0'][tmp]=[]
                        data_dict['S1'][tmp]=[]
                        data_dict['S2'][tmp]=[]
                        data_dict['DOLP'][tmp]=[]
                        
                    data_dict['S0'][tmp] += [os.path.join(root,domain,cat,'S0',s0path)]
                    data_dict['S1'][tmp] += [os.path.join(root,domain,cat,'S1',s1path)]
                    data_dict['S2'][tmp] += [os.path.join(root,domain,cat,'S2',s2path)]
                    data_dict['DOLP'][tmp] += [os.path.join(root,domain,cat,'DOLP',dolppath)]
                    
                # data_dict['S0'].append(os.path.join(root,domain,cat,'S0',s0path))
                # data_dict['S1'].append(os.path.join(root,domain,cat,'S0',s0path))
                # data_dict['S2'].append(os.path.join(root,domain,cat,'S0',s0path))
                # data_dict['DOLP'].append(os.path.join(root,domain,cat,'DOLP',dolppath))

                    
    traindir = {}
    testdir = {}
    traindir['S0'] = []
    traindir['S1'] = []
    traindir['S2'] = []
    traindir['DOLP'] = []
    
    testdir['S0'] = []
    testdir['S1'] = []
    testdir['S2'] = []
    testdir['DOLP'] = []

    # for modality in data_dict.keys():
    for cat in data_dict['S0'].keys():
        testlen = int(len(data_dict['S0'][cat]) *testratio)
        ind = np.arange(len(data_dict['S0'][cat]))
        random.shuffle(ind)
        testinds = ind[:testlen]
        traininds = ind[testlen:]
        for testind in testinds:
            testdir['S0'].append(data_dict['S0'][cat][testind])
            testdir['S1'].append(data_dict['S1'][cat][testind])
            testdir['S2'].append(data_dict['S2'][cat][testind])
            testdir['DOLP'].append(data_dict['DOLP'][cat][testind])
        
        for trainind in traininds:
            traindir['S0'].append(data_dict['S0'][cat][trainind])
            traindir['S1'].append(data_dict['S1'][cat][trainind])
            traindir['S2'].append(data_dict['S2'][cat][trainind])
            traindir['DOLP'].append(data_dict['DOLP'][cat][trainind])
                    
       
    return traindir,testdir

def initdirs(dirs):
    if os.path.exists(dirs[0]) is False:
        for _dir in dirs:
            os.makedirs(_dir)
    else:
        def del_dirsfile(dirs):
            for _dir in dirs:
                files = os.listdir(_dir)
                for file in files:
                    os.remove(os.path.join(_dir,file))
        del_dirsfile(dirs)

if __name__ == '__main__':
    datadir = os.path.join(os.path.abspath('.'),'datasets','CDFA')
    
    mode = 'gen1'
    
    if mode == 'gen1':
        ratio = 0.1
        version = 'v'+str(int(100*ratio))
        dis = '42_52_62'
        txtdir = os.path.join(os.path.abspath('.'),'data','V202207110920','polar_generation_'+dis)
        initdirs([txtdir])
        
        traindir,testdir = gendatalist(datadir,dis,ratio)
        print('trainLen:'+str(len(traindir['S0']))+'  testLen:'+str(len(testdir['S0'])))

        
        save_train_test_to_txt(np.array(traindir['S0']),np.array(testdir['S0']),txtdir,'S0_',version)
        save_train_test_to_txt(np.array(traindir['S1']),np.array(testdir['S1']),txtdir,'S1_',version)
        save_train_test_to_txt(np.array(traindir['S2']),np.array(testdir['S2']),txtdir,'S2_',version)
        save_train_test_to_txt(np.array(traindir['DOLP']),np.array(testdir['DOLP']),txtdir,'DOLP_',version)
    
    pass